package com.tianji.aigc.service.impl;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.map.MapUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.tianji.aigc.config.SystemPromptConfig;
import com.tianji.aigc.config.ToolResultHolder;
import com.tianji.aigc.constants.Constant;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.service.ChatService;
import com.tianji.aigc.service.ChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.utils.UserContext;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author zgp
 * @version V1.0
 * @date 2025-04-21 15:21
 */
// @Service
@RequiredArgsConstructor
public class ChatServiceImpl implements ChatService {
    private final VectorStore vectorStore;
    private final ChatClient dashScopeChatClient;
    private final SystemPromptConfig systemPromptConfig;
    private final ChatSessionService chatSessionService;
    // 对话输出状态标识
    private static final Map<String, Boolean> GENERATE_STATUS = new ConcurrentHashMap<>();
    private final ChatMemory chatMemory;
    // 输出结束的标记
    private static final ChatEventVO STOP_EVENT = ChatEventVO.builder().eventType(ChatEventTypeEnum.STOP.getValue()).build();

    /**
     * 流式对话聊天
     *
     * @param question  用户问题
     * @param sessionId 会话id
     * @return 大模型返回内容
     */
    @Override
    public Flux<ChatEventVO> chat(String question, String sessionId) {
        // 异步更新会话标题
        chatSessionService.updateSessionTitle(sessionId, UserContext.getUser(), question);
        // 生成请求id，用于保存tool执行结果时，作为唯一标识key
        String requestId = IdUtil.fastSimpleUUID();
        StringBuilder assistantMessageContentStr = new StringBuilder();
        // 获取对话id，用户id_会话id
        String conversationId = ChatService.getConversationId(sessionId);
        return dashScopeChatClient.prompt()
                .system(prompt -> prompt
                        // 设置系统提示词
                        .text(systemPromptConfig.getChatSystemMessage().get())
                        // 设置提示词中的参数
                        .param("now", DateUtil.now())
                )
                .user(question)
                .advisors(advisor -> advisor
                        .advisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.builder().query("").topK(999).build()))
                        .param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId))
                // 通过tool上下文，传递tool工具调用中所需参数
                .toolContext(Map.of(Constant.REQUEST_ID, requestId, Constant.USER_ID, UserContext.getUser()))
                .stream()
                .chatResponse()
                // 处理第一个流数据的时候，添加输出状态标识，true代表允许输出
                .doFirst(() -> GENERATE_STATUS.put(sessionId, true))
                // 处理最后一个流数据的时候，移除输出状态标识，代表输出结束
                .doOnComplete(() -> GENERATE_STATUS.remove(sessionId))
                // 处理异常的时候，移除输出状态标识，代表输出结束
                .doOnError(t -> GENERATE_STATUS.remove(sessionId))
                .doOnCancel(() -> {
                    // 会话停止生成时（即在FLux流取消时），持久化会话内容
                    saveStopChatAssistantContentToMemory(conversationId, assistantMessageContentStr.toString());
                })
                // 输出流数据时，判断输出状态标识，如果输出状态标识为true，则输出流数据，false则不输出流数据
                .takeWhile(c -> GENERATE_STATUS.getOrDefault(sessionId, false))
                .map(response -> {
                    String finishReason = response.getResult().getMetadata().getFinishReason();
                    if(StrUtil.equals(finishReason, Constant.FLUX_MESSAGE_STOP)){
                        // 最后一次流式结束时，将消息id和requestId进行关联
                        // 用于获取ChatEventTypeEnum.PARAM类型的数据，进行持久化
                        String messageId = response.getMetadata().getId();
                        ToolResultHolder.put(messageId, Constant.REQUEST_ID, requestId);
                    }
                    // 获取大模型的返回内容
                    String assistantMsg = response.getResult().getOutput().getText();
                    assistantMessageContentStr.append(assistantMsg);
                    // 包装内容为VO对象
                    return ChatEventVO.builder()
                            .eventData(assistantMsg)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                })
                .concatWith(Flux.defer(() -> {
                            // 1. 从全局容器中判断获取tool执行结果
                            Map<String, Object> toolExecuteResult = ToolResultHolder.get(requestId);
                            if (MapUtil.isNotEmpty(toolExecuteResult)) {
                                // 使用完清除tool执行结果
                                ToolResultHolder.remove(requestId);
                                // 2. 将tool执行结果添加到大模型返回内容中
                                return Flux.just(ChatEventVO.builder().eventType(ChatEventTypeEnum.PARAM.getValue()).eventData(toolExecuteResult).build(), STOP_EVENT);
                            }
                            return Flux.just(STOP_EVENT);
                        }
                ));
    }

    /**
     * 停止内容生成（中断流的响应，大模型内部依旧在生成内容）
     *
     * @param sessionId 会话id
     */
    @Override
    public void stopChat(String sessionId) {
        // 移除会话标识，代表停止内容输出
        GENERATE_STATUS.remove(sessionId);
    }

    /**
     * 持久化停止内容生成后的大模型返回内容
     *
     * @param conversationId   对话id
     * @param assistantContent 大模型的返回内容
     */
    private void saveStopChatAssistantContentToMemory(String conversationId, String assistantContent) {
        // 持久化大模型的返回内容到redis/mysql
        chatMemory.add(conversationId, List.of(new AssistantMessage(assistantContent)));
    }
}
